import neopixel
import machine
import time

class wm_rgb_ws2812:
    """WS2812 RGB LED灯带驱动类"""
    
    def __init__(self, pin=13, count=5):
        """
        初始化WS2812驱动
        
        Args:
            pin: GPIO引脚号，默认13
            count: LED数量，默认5
        """
        self.pin = machine.Pin(pin)
        self.rgb_count = count
        self.np = neopixel.NeoPixel(self.pin, count)
        
    def write(self, index, r, g, b):
        """
        设置指定LED的颜色
        
        Args:
            index: LED索引 (0=全部, 1-5=单个LED)
            r: 红色值 (0-255)
            g: 绿色值 (0-255)
            b: 蓝色值 (0-255)
        """
        if index == 0:
            # 设置所有LED
            for i in range(self.rgb_count):
                self.np[i] = (r, g, b)
        
        elif 1 <= index <= self.rgb_count:
            # 设置单个LED
            self.np[index - 1] = (r, g, b)
        
        else:
            raise ValueError(f"Index out of range. Must be 0-{self.rgb_count}")
        
        self.np.write()
    
    def clear(self):
        """清除所有LED（关闭）"""
        for i in range(self.rgb_count):
            self.np[i] = (0, 0, 0)
        self.np.write()
    
    def rainbow(self, delay=0.1):
        """彩虹效果"""
        colors = [
            (255, 0, 0),    # 红
            (255, 127, 0),  # 橙
            (255, 255, 0),  # 黄
            (0, 255, 0),    # 绿
            (0, 0, 255),    # 蓝
            (75, 0, 130),   # 靛
            (148, 0, 211)   # 紫
        ]
        
        for color in colors:
            self.write(0, *color)
            time.sleep(delay)
    
    def breathing(self, r, g, b, duration=2.0):
        """呼吸灯效果"""
        steps = 50
        for i in range(steps):
            # 渐亮
            brightness = i / steps
            self.write(0, int(r * brightness), int(g * brightness), int(b * brightness))
            time.sleep(duration / (steps * 2))
        
        for i in range(steps):
            # 渐暗
            brightness = (steps - i) / steps
            self.write(0, int(r * brightness), int(g * brightness), int(b * brightness))
            time.sleep(duration / (steps * 2))
    
    def chase(self, r, g, b, delay=0.1):
        """追逐效果"""
        for i in range(self.rgb_count):
            self.clear()
            self.write(i + 1, r, g, b)
            time.sleep(delay)
    
    def set_all_colors(self, colors):
        """
        设置所有LED为不同颜色
        
        Args:
            colors: 颜色列表，每个元素为(r,g,b)元组
        """
        if len(colors) != self.rgb_count:
            raise ValueError(f"Colors list length must be {self.rgb_count}")
        
        for i, color in enumerate(colors):
            self.np[i] = color
        self.np.write()